import os, tempfile

from pdbdataconnector import PDBDataConnector

class MMcifCategoryNotFoundException(Exception) : pass # can be used by any connector

class MMcifAPIconnector(PDBDataConnector) :
	'''this connector makes use of D&A api library provided by RCSB to read/write mmcif.  only the first datablock is used!'''
	def __init__(s, filepathORurl) :
		PDBDataConnector()
		from data_access.CifFile import CifFile ## see README for setup necessary to use D&A api
		tempfd, tempfn = tempfile.mkstemp(suffix=".txt",prefix="MMcifAPIconnectorRcsbLog") ; os.close(tempfd)
		s.datablock = CifFile(s.makeTempFile(filepathORurl), tempfn).getCifFile()
		s.datablock = s.datablock.GetBlock( s.datablock.GetFirstBlockName() )
	def getCategory(s, cat, item=None, mode="normal") :
		'''overriding the parent method'''
		mult = 0
		try : table = s.datablock.GetTable(cat)
		except : raise MMcifCategoryNotFoundException("No category found: " + cat)
		mult = table.GetNumRows()
		if mode == "normal" :
			if mult == 1 : return MMcifAPIconnector.catRetMask(table, 0, item)
			else : return MMcifAPIconnector.iterCat(table, item)
		elif mode == "iter" :
			return MMcifAPIconnector.iterCat(table, item)
		elif mode == "first" :
			return MMcifAPIconnector.catRetMask(table, 0, item)
		else : raise NotImplementedError
	@staticmethod
	def iterCat(table, item) :
		'''iterator loop that yields a map or item's value'''
		for ri in range(table.GetNumRows()) :
			yield MMcifAPIconnector.catRetMask(table, ri, item)
	@staticmethod
	def catRetMask(table, ri, item) :
		'''either return dictionary or just item's value for ri'th row in table'''
		if item == None : return MMcifAPIconnector.itemMap(table, ri)
		else :
			try : return table(ri, item)
			except : raise MMcifCategoryNotFoundException("No category/item found: " + table.GetName() + " " + item)
	@staticmethod
	def itemMap(table, ri) :
		'''make a dictionary from column names to values for ri'th row in table'''
		retmap = {}
		for cn in table.GetColumnNames() : retmap[cn] = table(ri,cn)
		return retmap
